import os
import tools
import json
import pickle

def createIndexTable(directname,indexname):
    invertedIndex = {}      # 倒排索引表
    path = tools.projectpath + directname       # 存放分词后的语料集
    files = os.listdir(path)

    indexpath = tools.projectpath + indexname     # 存放索引的文件名
    wordlist = []
    indexfile = open(indexpath, 'rb')
    wordlist.extend(pickle.load(indexfile))
    print(wordlist)

    # 建立倒排索引表
    for file in files:
        print("analyzing file: ", file)
        temp = open(path + '/' + file,'rb')
        content = pickle.load(temp)["context"]
        # content = pickle.load(temp)["title_cut"]

        docId = getDocID(file)
        print(docId)

        num = 0         # word在文档中的位置
        for word in content:
            if word in wordlist:
                print(word)
                addIndex(word, invertedIndex, docId, num)

            num += 1
    # 给倒排索引中的词项排序
    invertedIndex = sortTheDict(invertedIndex)

    printIndex(invertedIndex)

    # 将数据写入文件中
    with open(tools.projectpath + 'out/invertIndex.json', "w", encoding='utf-8') as f:
        json.dump(invertedIndex, f, ensure_ascii=False)
        print(u'加载入文件完成...')

# 增加倒排索引项
def addIndex(word,invertedIndex,docId,num):
    if word not in invertedIndex:
        docList = {}
        docList[docId] = [num]
        invertedIndex[word] = docList
    else:
        if docId not in invertedIndex[word]:
            invertedIndex[word][docId] = [num]
        else:
            invertedIndex[word][docId].append(num)

#获取文档名中的文档的id
def getDocID(filename):
    # end = filename.find('.')
    docId = filename[:]
    return int(docId)

def sortTheDict(dict):
    sdict =  { k:dict[k] for k in sorted(dict.keys())}
    for stem in sdict:
        sdict[stem] = { k:sdict[stem][k] for k in sorted(sdict[stem].keys())}
    return sdict

def printIndex(index):
    for stem in index:
        print(stem)
        for doc in index[stem]:
            print("    " , doc , " : " , index[stem][doc])

# 测试
createIndexTable("data/data_clean_key_cut","data/all_key_words_list_add_syn_extend")
# createIndexTable("data/data_clean_key_cut","data/key_words_title")